"""
任务队列管理路由
基于asyncio实现的AI分析任务队列系统
提供任务创建、状态查询、结果获取等功能
"""

from typing import Dict, Optional
import asyncio
import uuid
from datetime import datetime
from fastapi import APIRouter
from loguru import logger

from models.schemas import AnalysisRequest, TaskCreateResponse, TaskResultResponse, StatusCode, TaskStatus, TaskListResponse
from routes.ai_routes import request_ai_analysis

# 创建任务路由器
task_router = APIRouter()

# =====================
# 任务队列核心组件
# =====================

# 任务队列：存储待处理的任务
# 队列元素格式：(task_id: str, request: AnalysisRequest)
_task_queue: asyncio.Queue = asyncio.Queue()

# 任务状态存储：内存字典，存储所有任务的状态信息
# 格式：task_id -> {
#   status: str,           # 任务状态：pending/processing/completed/failed
#   created_at: str,       # 任务创建时间戳
#   finished_at: str,      # 任务完成时间戳
#   request: dict,         # 原始请求数据
#   result: dict,          # 任务执行结果（成功时）
#   error: str             # 错误信息（失败时）
# }
_task_storage: Dict[str, dict] = {}

# 配置：最大保存的任务数量（防止内存溢出）
MAX_TASKS_IN_MEMORY = 100

# 配置：任务保留时间（秒），超过此时间的已完成任务会被自动清理
TASK_RETENTION_SECONDS = 3600  # 1小时

# 后台工作器任务句柄
# 用于控制和管理后台任务处理器的生命周期
_worker_task: Optional[asyncio.Task] = None

# =====================
# 任务清理功能
# =====================

def _clean_old_tasks() -> int:
    """
    清理旧任务，防止内存溢出
    
    清理策略：
    1. 如果任务总数超过 MAX_TASKS_IN_MEMORY，删除最老的已完成/失败任务
    2. 删除超过 TASK_RETENTION_SECONDS 的已完成/失败任务
    3. 不删除 pending 和 running 状态的任务
    
    返回：清理的任务数量
    """
    from datetime import datetime, timedelta
    
    current_time = datetime.now()
    tasks_to_delete = []
    
    # 策略1: 如果任务总数超限，按时间排序删除最老的已完成任务
    if len(_task_storage) > MAX_TASKS_IN_MEMORY:
        # 筛选出已完成和失败的任务
        finished_tasks = [
            (task_id, task_meta.get("finished_at", ""))
            for task_id, task_meta in _task_storage.items()
            if task_meta.get("status") in [TaskStatus.COMPLETED.value, TaskStatus.FAILED.value]
        ]
        
        # 按完成时间排序（最老的在前）
        finished_tasks.sort(key=lambda x: x[1])
        
        # 计算需要删除多少个任务
        num_to_delete = len(_task_storage) - MAX_TASKS_IN_MEMORY + 100  # 多删除100个，避免频繁清理
        tasks_to_delete.extend([task_id for task_id, _ in finished_tasks[:num_to_delete]])
    
    # 策略2: 删除超时的已完成任务
    for task_id, task_meta in _task_storage.items():
        if task_id in tasks_to_delete:
            continue
            
        status = task_meta.get("status")
        finished_at_str = task_meta.get("finished_at")
        
        # 只清理已完成或失败的任务
        if status in [TaskStatus.COMPLETED.value, TaskStatus.FAILED.value] and finished_at_str:
            try:
                finished_at = datetime.strptime(finished_at_str, "%Y-%m-%d %H:%M:%S")
                age_seconds = (current_time - finished_at).total_seconds()
                
                if age_seconds > TASK_RETENTION_SECONDS:
                    tasks_to_delete.append(task_id)
            except Exception as e:
                logger.warning(f"解析任务 {task_id} 完成时间失败: {e}")
    
    # 执行删除
    deleted_count = 0
    for task_id in tasks_to_delete:
        if task_id in _task_storage:
            # 删除前释放请求数据中的大对象
            task_meta = _task_storage[task_id]
            if "request" in task_meta:
                task_meta["request"] = None  # 释放原始请求数据
            del _task_storage[task_id]
            deleted_count += 1
    
    if deleted_count > 0:
        logger.info(f"🧹 清理了 {deleted_count} 个旧任务，当前任务总数: {len(_task_storage)}")
    
    return deleted_count


# =====================
# 任务处理核心逻辑
# =====================

async def _task_worker() -> None:
    """
    后台任务处理器（Worker）
    
    功能说明：
    1. 持续监听任务队列，获取待处理任务
    2. 更新任务状态为"处理中"
    3. 调用AI分析接口执行具体任务
    4. 保存执行结果或错误信息
    5. 更新任务状态为"已完成"或"失败"
    6. 定期清理旧任务，防止内存溢出
    
    注意事项：
    - 单线程顺序处理，避免资源竞争
    - 异常处理确保单个任务失败不影响整体队列
    - 使用asyncio确保非阻塞处理
    """
    logger.info("任务处理器启动，开始监听任务队列...")
    
    # 任务计数器，用于定期清理
    task_count = 0
    
    while True:
        try:
            # 从队列中获取任务（阻塞等待）
            task_id, analysis_request = await _task_queue.get()
            logger.info(f"开始处理任务: {task_id}")
            
            # 获取任务元信息并更新状态
            task_meta = _task_storage.get(task_id, {})
            task_meta["status"] = TaskStatus.RUNNING.value
            _task_storage[task_id] = task_meta
            
            try:
                # 执行AI分析任务
                analysis_result = await request_ai_analysis(analysis_request)
                
                # 任务执行成功，保存结果
                task_meta["status"] = TaskStatus.COMPLETED.value
                task_meta["finished_at"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                task_meta["processing_time"] = analysis_result.processing_time
                task_meta["dicom_uid"] = analysis_request.dicom_uid
                task_meta["result"] = analysis_result.result  # 转换为字典格式存储
                
                logger.info(f"任务 {task_id} 执行成功")
                
            except Exception as task_error:
                # 任务执行失败，记录详细错误信息
                logger.error(f"任务 {task_id} 执行失败: {str(task_error)}")
                
                # 记录完整的异常堆栈跟踪信息
                logger.exception(f"任务 {task_id} 详细错误信息:")
                
                # 记录任务上下文信息用于调试
                logger.error(f"任务 {task_id} 请求参数: dicom_uid={analysis_request.dicom_uid}, "
                           f"dsa_position={analysis_request.dsa_position}")
                
                # 记录任务处理时间信息
                created_time = task_meta.get("created_at", "未知")
                logger.error(f"任务 {task_id} 处理时间信息: 创建时间={created_time}, "
                           f"失败时间={datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
                
                task_meta["status"] = TaskStatus.FAILED.value
                task_meta["finished_at"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                task_meta["error"] = str(task_error)
                raise task_meta
            
            finally:
                # 更新任务存储并标记队列任务完成
                _task_storage[task_id] = task_meta
                
                # 增加任务计数器
                task_count += 1
                
                # 每处理10个任务清理一次旧任务
                if task_count % 10 == 0:
                    _clean_old_tasks()
                
        except Exception as worker_error:
            # Worker级别的异常处理，确保Worker不会崩溃
            logger.error(f"任务处理器发生异常: {str(worker_error)}")
            await asyncio.sleep(1)  # 短暂休眠后继续处理

async def _ensure_worker_running() -> None:
    """
    确保后台任务处理器正在运行
    
    功能说明：
    - 检查Worker是否已经启动
    - 如果未启动或已结束，则创建新的Worker任务
    - 采用懒加载方式，只在需要时启动Worker
    
    设计考虑：
    - 不能在模块导入时启动Worker，因为此时事件循环可能未就绪
    - 在第一次创建任务时调用此函数启动Worker
    - 支持Worker异常退出后的自动重启
    """
    global _worker_task
    
    # 检查Worker是否需要启动
    if _worker_task is None or _worker_task.done():
        logger.info("启动后台任务处理器...")
        loop = asyncio.get_running_loop()
        _worker_task = loop.create_task(_task_worker())

# =====================
# 任务管理API接口
# =====================

@task_router.post("/tasks", summary="创建AI分析任务", response_model=TaskCreateResponse)
async def create_task(request: AnalysisRequest):
    """
    创建新的AI分析任务
    
    请求参数：
    - request: AnalysisRequest - AI分析请求对象，包含dicom_uid、img_data、seg_data、dsa_position等
    
    返回结果：
    - task_id: str - 唯一任务标识符
    - status: str - 任务初始状态（"pending"）
    - created_at: float - 任务创建时间戳
    
    使用示例：
    POST /tasks
    {
        "dicom_uid": "sample_uid_123",
        "img_data": [[0,0,0],[0,0,0],[0,0,0]],
        "seg_data": [[0,0,0],[0,0,0],[0,0,0]],
        "dsa_position": "F"
    }
    
    响应示例：
    {
        "task_id": "a1b2c3d4e5f6...",
        "status": "pending",
        "created_at": 1640995200.123
    }
    """
    # 确保后台处理器正在运行
    await _ensure_worker_running()
    
    # 生成唯一任务标识符
    task_id = uuid.uuid4().hex
    # 时间改成年月日
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # 在创建新任务前清理旧任务，防止内存溢出
    _clean_old_tasks()
    
    # 初始化任务状态信息
    task_meta = {
        "status": TaskStatus.PENDING.value,
        "created_at": current_time,
        "finished_at": None,
        "dicom_uid": request.dicom_uid,  # 添加dicom_uid字段
        "result": None,
        "error": None,
        "request": request  # 保存原始请求数据
    }
    
    # 存储任务信息
    _task_storage[task_id] = task_meta
    
    # 将任务加入处理队列
    await _task_queue.put((task_id, request))
    
    logger.info(f"创建新任务: {task_id}, DICOM UID: {request.dicom_uid}")
    
    return TaskCreateResponse(
        task_id=task_id,
        status=TaskStatus.PENDING,
        created_at=current_time,
        status_code=StatusCode.SUCCESS
    )

@task_router.get("/tasks/{task_id}", summary="查询任务状态和结果", response_model=TaskResultResponse)
async def get_task_status(task_id: str):
    """
    根据任务ID查询任务状态和结果
    
    路径参数：
    - task_id: str - 任务唯一标识符
    
    返回结果：
    根据任务状态返回不同内容：
    
    基础信息（所有状态都包含）：
    - task_id: str - 任务ID
    - status: str - 任务状态
    - created_at: str - 创建时间
    
    成功完成时额外包含：
    - result: dict - AI分析结果
    
    失败时额外包含：
    - error: str - 错误信息
    
    任务状态说明：
    - pending: 等待处理
    - processing: 正在处理
    - completed: 处理完成
    - failed: 处理失败

    """
    # 查询任务信息
    task_meta = _task_storage.get(task_id)
    
    # 检查任务是否存在
    if task_meta is None:
        logger.warning(f"查询不存在的任务: {task_id}")
        return TaskResultResponse(
            task_id=task_id,
            status=TaskStatus.FAILED,
            created_at=None,
            finished_at=None,
            result=None,
            error=f"任务不存在: {task_id}",
            status_code=StatusCode.NOT_FOUND.value
        )
    
    # 构建基础响应数据
    task_status = task_meta.get("status")
    
    # 根据任务状态确定状态码和状态信息
    if task_status == TaskStatus.COMPLETED.value:
        status_code = StatusCode.SUCCESS
        status_enum = TaskStatus.COMPLETED
    elif task_status == TaskStatus.FAILED.value:
        status_code = StatusCode.INTERNAL_SERVER_ERROR
        status_enum = TaskStatus.FAILED
    elif task_status == TaskStatus.RUNNING.value:
        status_code = StatusCode.SUCCESS
        status_enum = TaskStatus.RUNNING
    elif task_status == TaskStatus.PENDING.value:
        status_code = StatusCode.SUCCESS
        status_enum = TaskStatus.PENDING
    else:
        status_code = StatusCode.SUCCESS
        status_enum = TaskStatus.PENDING

    response_data = TaskResultResponse(
            task_id=task_id,
            status=status_enum,
            created_at=task_meta.get("created_at"),
            finished_at=task_meta.get("finished_at"),
            dicom_uid=task_meta.get("dicom_uid"),
            result=task_meta.get("result"),
            error=task_meta.get("error"),
            status_code=status_code
        )
    
    return response_data

# 获取任务列表
@task_router.post("/tasks/list", summary="查询任务列表", response_model=TaskListResponse)
async def get_task_list():  
    """
    获取任务列表
    """
    task_list = []
    for task_id, task_meta in _task_storage.items():
        item = {
            "task_id": task_id,
            "status": task_meta.get("status"),
            "created_at": task_meta.get("created_at"),
            "finished_at": task_meta.get("finished_at"),
            "dicom_uid": task_meta.get("dicom_uid"),
            "result": task_meta.get("result"),
            "error": task_meta.get("error")
        }
        task_list.append(item)
    return TaskListResponse(
        task_list=task_list,
        status_code=StatusCode.SUCCESS
    )
    